#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""  
@Project : PythonProject
@File : database.py.py
@Author : donc
@Time : 2025/5/28 14:08  
@脚本说明 : 
    
"""

import streamlit as st
import logging
import hashlib
import base64
import pymysql
import urllib.parse
from passlib.hash import pbkdf2_sha256  # 添加缺失的导入


# 配置日志
logger = logging.getLogger(__name__)
def get_db_connection():
    """获取 MySQL 数据库连接"""
    try:
        conn = pymysql.connect(
            host=st.secrets.mysql.host,
            port=st.secrets.mysql.port,
            user=st.secrets.mysql.user,
            password=st.secrets.mysql.password,
            database=st.secrets.mysql.database,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )
        logger.info("数据库连接成功")
        return conn
    except pymysql.Error as e:
        logger.error(f"数据库连接失败: {e}")
        st.error(f"数据库连接失败: {e}")
        return None


def get_db_connection():
    """获取 MySQL 数据库连接"""
    try:
        conn = pymysql.connect(
            host=st.secrets.mysql.host,
            port=st.secrets.mysql.port,
            user=st.secrets.mysql.user,
            password=st.secrets.mysql.password,
            database=st.secrets.mysql.database,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )
        logger.info("数据库连接成功")
        return conn
    except pymysql.Error as e:
        logger.error(f"数据库连接失败: {e}")
        st.error(f"数据库连接失败: {e}")
        return None


def init_db():
    """初始化数据库和用户表"""
    conn = get_db_connection()
    if not conn:
        return False

    try:
        with conn.cursor() as cursor:
            # 创建用户表
            cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INT AUTO_INCREMENT PRIMARY KEY,
                username VARCHAR(50) UNIQUE NOT NULL,
                password VARCHAR(255) NOT NULL,
                email VARCHAR(100),
                role VARCHAR(20) DEFAULT 'user',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_login TIMESTAMP NULL
            )
            ''')

            # 创建管理员用户（如果不存在）
            admin_password = pbkdf2_sha256.hash("admin123")
            cursor.execute("SELECT * FROM users WHERE username='admin'")
            if not cursor.fetchone():
                cursor.execute(
                    "INSERT INTO users (username, password, role) VALUES (%s, %s, %s)",
                    ('admin', admin_password, 'admin')
                )
                logger.info("管理员用户创建成功")

            # 创建历史记录表
            cursor.execute('''
            CREATE TABLE IF NOT EXISTS history (
                id INT AUTO_INCREMENT PRIMARY KEY,
                input_text TEXT NOT NULL,
                output_text TEXT NOT NULL,
                method VARCHAR(50) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            ''')

            # 创建 MD5 字典表
            cursor.execute('''
            CREATE TABLE IF NOT EXISTS md5_dictionary (
                id INT AUTO_INCREMENT PRIMARY KEY,
                encrypted VARCHAR(255) NOT NULL,
                original TEXT NOT NULL
            )
            ''')

            # 创建 MySQL 字典表
            cursor.execute('''
            CREATE TABLE IF NOT EXISTS mysql_dictionary (
                id INT AUTO_INCREMENT PRIMARY KEY,
                encrypted VARCHAR(255) NOT NULL,
                original TEXT NOT NULL
            )
            ''')

            # 创建SHA系列字典表
            cursor.execute('''
            CREATE TABLE IF NOT EXISTS sha1_dictionary (
                id INT AUTO_INCREMENT PRIMARY KEY,
                encrypted VARCHAR(40) NOT NULL,
                original TEXT NOT NULL
            )
            ''')

            cursor.execute('''
            CREATE TABLE IF NOT EXISTS sha256_dictionary (
                id INT AUTO_INCREMENT PRIMARY KEY,
                encrypted VARCHAR(64) NOT NULL,
                original TEXT NOT NULL
            )
            ''')

            cursor.execute('''
            CREATE TABLE IF NOT EXISTS sha512_dictionary (
                id INT AUTO_INCREMENT PRIMARY KEY,
                encrypted VARCHAR(128) NOT NULL,
                original TEXT NOT NULL
            )
            ''')

            # 创建 RailFence 字典表
            cursor.execute('''
            CREATE TABLE IF NOT EXISTS railfence_dictionary (
                id INT AUTO_INCREMENT PRIMARY KEY,
                encrypted TEXT NOT NULL,
                original TEXT NOT NULL
            )
            ''')

            conn.commit()
            return True
    except pymysql.Error as e:
        logger.error(f"数据库初始化失败: {e}")
        st.error(f"数据库初始化失败: {e}")
        return False
    finally:
        if conn:
            conn.close()


def get_user(username):
    """根据用户名获取用户信息"""
    conn = get_db_connection()
    if not conn:
        return None

    try:
        with conn.cursor() as cursor:
            cursor.execute("SELECT * FROM users WHERE username = %s", (username,))
            return cursor.fetchone()
    except pymysql.Error as e:
        logger.error(f"获取用户信息失败: {e}")
        return None
    finally:
        conn.close()


def create_user(username, password, email, role="user"):
    """创建新用户"""
    conn = get_db_connection()
    if not conn:
        return False

    try:
        hashed_password = pbkdf2_sha256.hash(password)
        with conn.cursor() as cursor:
            cursor.execute(
                "INSERT INTO users (username, password, email, role) VALUES (%s, %s, %s, %s)",
                (username, hashed_password, email, role)
            )
            conn.commit()
            return cursor.rowcount > 0  # 返回是否成功插入
    except pymysql.IntegrityError as e:
            # 处理唯一约束冲突（用户名已存在）
        logger.error(f"创建用户失败（用户名已存在）: {e}")
        return False
    except pymysql.Error as e:
        logger.error(f"创建用户失败: {e}")
        return False
    finally:
        conn.close()


def update_password(username, new_password):
    """更新用户密码"""
    conn = get_db_connection()
    if not conn:
        return False

    try:
        hashed_password = pbkdf2_sha256.hash(new_password)
        with conn.cursor() as cursor:
            cursor.execute(
                "UPDATE users SET password = %s WHERE username = %s",
                (hashed_password, username)
            )
            conn.commit()
            return cursor.rowcount > 0
    except pymysql.Error as e:
        logger.error(f"更新密码失败: {e}")
        return False
    finally:
        conn.close()

def get_user_role(username):
    """获取用户角色"""
    conn = get_db_connection()
    if not conn:
        return None

    try:
        with conn.cursor() as cursor:
            cursor.execute("SELECT role FROM users WHERE username = %s", (username,))
            result = cursor.fetchone()
            return result['role'] if result else 'user'
    except pymysql.Error as e:
        logger.error(f"获取用户角色失败: {e}")
        return 'user'
    finally:
        if conn:
            conn.close()


def get_all_users():
    """获取所有用户信息"""
    conn = get_db_connection()
    if not conn:
        return []

    try:
        with conn.cursor() as cursor:
            cursor.execute("""
                            SELECT id, username, email, role, 
                                   DATE_FORMAT(created_at, '%Y-%m-%d %H:%i') AS created_at,
                                   IFNULL(DATE_FORMAT(last_login, '%Y-%m-%d %H:%i'), '-') AS last_login 
                            FROM users
                        """)
            return cursor.fetchall()
    except pymysql.Error as e:
        logger.error(f"获取用户列表失败: {e}")
        return []
    finally:
        if conn:
            conn.close()


def update_last_login(username):
    """更新最后登录时间"""
    conn = get_db_connection()
    if not conn:
        return False

    try:
        with conn.cursor() as cursor:
            cursor.execute(
                "UPDATE users SET last_login = CURRENT_TIMESTAMP WHERE username = %s",
                (username,)
            )
            conn.commit()
            return True
    except pymysql.Error as e:
        logger.error(f"更新登录时间失败: {e}")
        return False
    finally:
        if conn:
            conn.close()

def verify_password(password, hashed_password):
    """验证密码"""
    try:
        return pbkdf2_sha256.verify(password, hashed_password)
    except Exception as e:
        logger.error(f"密码验证失败: {e}")
        return False


def delete_user(user_id):
    """删除用户"""
    conn = get_db_connection()
    if not conn:
        return False

    try:
        with conn.cursor() as cursor:
            cursor.execute("DELETE FROM users WHERE id = %s", (user_id,))
            conn.commit()
            return cursor.rowcount > 0
    except pymysql.Error as e:
        logger.error(f"删除用户失败: {e}")
        return False
    finally:
        conn.close()


def admin_panel():
    """管理员面板 - 增强版用户管理"""

    st.subheader("👥 用户管理系统")

    # 显示所有用户
    users = get_all_users()

    # 使用选项卡布局
    add_tab, edit_tab, delete_tab = st.tabs(["➕ 添加用户", "✏️ 编辑用户", "🗑️ 删除用户"])
    with add_tab:
        st.subheader("添加新用户")
        with st.form("add_user_form", clear_on_submit=True):
            col1, col2 = st.columns(2)
            with col1:
                new_username = st.text_input("用户名*", key="new_username")
            with col2:
                new_email = st.text_input("邮箱", key="new_email")

            col3, col4 = st.columns(2)
            with col3:
                new_password = st.text_input("密码*", type="password", key="new_password")
            with col4:
                new_role = st.selectbox("角色*", ["user", "admin"], key="new_role")

            st.caption("带 * 的字段为必填项")
            add_user = st.form_submit_button("添加用户")

            if add_user:
                if not new_username or not new_password:
                    st.error("用户名和密码是必填项")
                    return

                # 检查用户名是否已存在
                if get_user(new_username):
                    st.error("用户名已存在")
                    return

                # 创建新用户
                if create_user(new_username, new_password, new_email, new_role):
                    st.success(f"用户 {new_username} 添加成功")
                    # 重新加载页面以刷新用户列表
                    st.rerun()
                else:
                    st.error("添加用户失败，请检查日志")

    with edit_tab:
        st.subheader("编辑用户信息")
        if not users:
            st.warning("没有可编辑的用户")
        else:
            # 创建用户选择器
            user_options = {f"{user['id']}-{user['username']}": user for user in users}
            selected_user = st.selectbox(
                "选择要编辑的用户",
                list(user_options.keys()),
                key="edit_user_select"
            )

            # 获取选中用户的信息
            user_data = user_options[selected_user]

            with st.form("edit_user_form"):
                st.info(f"当前编辑用户: {user_data['username']} (ID: {user_data['id']})")

                # 使用列布局
                col1, col2 = st.columns(2)
                with col1:
                    new_username = st.text_input("新用户名*", value=user_data['username'])
                with col2:
                    new_email = st.text_input("新邮箱", value=user_data.get('email', ''))

                col3, col4 = st.columns(2)
                with col3:
                    new_role = st.selectbox(
                        "新角色*",
                        ["user", "admin"],
                        index=0 if user_data['role'] == 'user' else 1
                    )
                with col4:
                    new_password = st.text_input("新密码（留空则不修改）", type="password",
                                                 help="如果不需要修改密码，请留空此字段")

                st.caption("带 * 的字段为必填项")
                update_user = st.form_submit_button("更新用户信息")

                if update_user:
                    # 检查用户名是否冲突（排除当前用户）
                    if new_username != user_data['username'] and get_user(new_username):
                        st.error("新用户名已被占用")
                        return

                    # 更新用户信息
                    conn = get_db_connection()
                    if not conn:
                        st.error("数据库连接失败")
                        return

                    try:
                        with conn.cursor() as cursor:
                            # 更新基本字段
                            update_query = """
                                UPDATE users 
                                SET username = %s, email = %s, role = %s 
                                WHERE id = %s
                            """
                            cursor.execute(
                                update_query,
                                (new_username, new_email, new_role, user_data['id'])
                            )

                            # 如果提供了新密码，则更新密码
                            if new_password:
                                hashed_password = pbkdf2_sha256.hash(new_password)
                                cursor.execute(
                                    "UPDATE users SET password = %s WHERE id = %s",
                                    (hashed_password, user_data['id'])
                                )

                            conn.commit()

                            # 如果更新了当前登录用户，更新会话状态
                            current_user_id = st.session_state.user.get("id")
                            if user_data['id'] == current_user_id:
                                st.session_state.user['username'] = new_username
                                st.session_state.user['email'] = new_email
                                st.session_state.user['role'] = new_role
                                st.session_state.LOGGED_IN_USERNAME = new_username
                                st.session_state.LOGGED_IN_ROLE = new_role

                            st.success("用户信息更新成功")
                            st.rerun()
                    except pymysql.Error as e:
                        logger.error(f"更新用户失败: {e}")
                        st.error(f"更新用户失败: {e}")
                    finally:
                        conn.close()

    with delete_tab:
        st.subheader("删除用户")
        if not users:
            st.warning("没有可删除的用户")
        else:
            # 使用两列布局
            col1, col2 = st.columns([1, 2])

            with col1:
                # 创建用户选择器
                user_options = {f"{user['id']}-{user['username']}": user for user in users}
                selected_user = st.selectbox(
                    "选择要删除的用户",
                    list(user_options.keys()),
                    key="delete_user_select"
                )

                # 获取选中用户的信息
                user_data = user_options[selected_user]
                user_id = user_data['id']

                # 显示警告信息
                current_user_id = st.session_state.user.get("id")
                if user_id == current_user_id:
                    st.error("⚠️ 不能删除当前登录的用户")
                else:
                    st.warning(f"确定要删除用户 {user_data['username']} 吗？此操作不可恢复！")

            with col2:
                with st.form("delete_user_form"):
                    # 添加确认输入框
                    confirm_text = st.text_input(
                        f"请输入 'DELETE {user_data['username']}' 确认删除",
                        key=f"confirm_delete_{user_id}"
                    )

                    # 删除按钮
                    delete_button = st.form_submit_button(
                        "确认删除用户",
                        disabled=(user_id == current_user_id)
                    )

                    if delete_button:
                        # 检查确认文本
                        if confirm_text != f"DELETE {user_data['username']}":
                            st.error("确认文本输入不正确")
                            return

                        if delete_user(user_id):
                            st.success(f"用户 {user_data['username']} 已删除")
                            st.rerun()
                        else:
                            st.error("删除用户失败")
def save_to_history(input_text, output_text, method):
    """保存操作历史记录"""
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            # 正确的字段名称
            sql = """
            INSERT INTO history 
            (input_text, output_text, method) 
            VALUES (%s, %s, %s)
            """
            cursor.execute(sql, (input_text, output_text, method))
        conn.commit()
    except Exception as e:
        logger.error(f"保存历史记录失败: {e}")
    finally:
        if conn:
            conn.close()

def get_encryption_history():
    """获取加密解密历史记录"""
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            # 使用正确的字段名称和别名
            sql = """
            SELECT 
                id,
                input_text AS input,
                output_text AS output,
                method,
                DATE_FORMAT(created_at, '%Y-%m-%d %H:%i:%s') AS time
            FROM history
            ORDER BY created_at DESC
            LIMIT 50
            """
            cursor.execute(sql)
            return cursor.fetchall()  # 直接返回字典列表
    except Exception as e:
        logger.error(f"获取历史记录失败: {e}")
        return []
    finally:
        if conn:
            conn.close()

def clear_history():
    """清除所有历史记录"""
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            cursor.execute("DELETE FROM history")
            conn.commit()
            return True
    except Exception as e:
        logger.error(f"清除历史记录失败: {e}")
        return False
    finally:
        if conn:
            conn.close()


def save_to_dictionary(encrypted, original, method):
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            table_name = f"{method.lower()}_dictionary"
            sql = f"INSERT INTO {table_name} (encrypted, original) VALUES (%s, %s)"
            cursor.execute(sql, (encrypted, original))
        conn.commit()
    finally:
        conn.close()


# ASCII 码转换功能
def text_to_ascii(text):
    ascii_values = [str(ord(char)) for char in text]
    result = ' '.join(ascii_values)
    save_to_history(text, result, 'ASCII')
    return result


def ascii_to_text(ascii_text):
    try:
        ascii_values = ascii_text.split()
        text = ''.join([chr(int(value)) for value in ascii_values])
        save_to_history(ascii_text, text, 'ASCII')
        return text
    except:
        return None


# MD5 加密解密
def md5_encrypt(text, salt=''):
    text = str(text) + str(salt)
    hash_object = hashlib.md5(text.encode())
    encrypted = hash_object.hexdigest()
    save_to_history(text, encrypted, 'MD5')
    save_to_dictionary(encrypted, text, 'MD5')
    return encrypted


def md5_decrypt(hash_str, salt=None):
    # 先查询历史记录
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            sql = "SELECT input_text FROM history WHERE output_text = %s AND method = 'MD5'"
            cursor.execute(sql, (hash_str,))
            result = cursor.fetchone()
            if result:
                original_text = result['input_text']
                if salt:
                    original_text = original_text.replace(salt, '')
                save_to_history(hash_str, original_text, 'MD5')
                return original_text
    finally:
        conn.close()

    # 再查询字典表
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            sql = "SELECT original FROM md5_dictionary WHERE encrypted = %s"
            cursor.execute(sql, (hash_str,))
            result = cursor.fetchone()
            if result:
                original_text = result['original']
                if salt:
                    original_text = original_text.replace(salt, '')
                save_to_history(hash_str, original_text, 'MD5')
                return original_text
            return None
    finally:
        conn.close()


# Base64 加密解密
def base64_encrypt(text):
    text = str(text).encode()
    encrypted = base64.b64encode(text).decode()
    save_to_history(text.decode(), encrypted, 'Base64')
    return encrypted


def base64_decrypt(encrypted_text):
    try:
        decrypted = base64.b64decode(encrypted_text).decode()
        save_to_history(encrypted_text, decrypted, 'Base64')
        return decrypted
    except:
        return None


# MySQL 加密解密
def mysql_encrypt(text, salt=''):
    text = str(text) + str(salt)
    hash_object = hashlib.sha256(text.encode())
    encrypted = '*' + hash_object.hexdigest().upper()
    save_to_history(text, encrypted, 'MySQL')
    save_to_dictionary(encrypted, text, 'MySQL')
    return encrypted


def mysql_decrypt(hash_str, salt=None):
    # 先查询历史记录
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            sql = "SELECT input_text FROM history WHERE output_text = %s AND method = 'MySQL'"
            cursor.execute(sql, (hash_str,))
            result = cursor.fetchone()
            if result:
                original_text = result['input_text']
                if salt:
                    original_text = original_text.replace(salt, '')
                save_to_history(hash_str, original_text, 'MySQL')
                return original_text
    finally:
        conn.close()

    # 再查询字典表
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            sql = "SELECT original FROM mysql_dictionary WHERE encrypted = %s"
            cursor.execute(sql, (hash_str,))
            result = cursor.fetchone()
            if result:
                original_text = result['original']
                if salt:
                    original_text = original_text.replace(salt, '')
                save_to_history(hash_str, original_text, 'MySQL')
                return original_text
            return None
    finally:
        conn.close()


# SHA 系列加密解密
def sha_encrypt(text, salt='', sha_type='sha256'):
    text = str(text) + str(salt)
    if sha_type == 'sha1':
        hash_object = hashlib.sha1(text.encode())
    elif sha_type == 'sha256':
        hash_object = hashlib.sha256(text.encode())
    elif sha_type == 'sha512':
        hash_object = hashlib.sha512(text.encode())
    else:
        return None
    encrypted = hash_object.hexdigest()
    save_to_history(text, encrypted, sha_type.upper())
    save_to_dictionary(encrypted, text, sha_type.upper())
    return encrypted


def sha_decrypt(hash_str, salt=None, sha_type='sha256'):
    # 先查询历史记录
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            sql = "SELECT input_text FROM history WHERE output_text = %s AND method = %s"
            cursor.execute(sql, (hash_str, sha_type.upper()))
            result = cursor.fetchone()
            if result is not None:  # 检查 result 是否为 None
                original_text = result['input_text']
                if salt:
                    original_text = original_text.replace(salt, '')
                save_to_history(hash_str, original_text, sha_type.upper())
                return original_text
    finally:
        conn.close()

    # 再查询字典表
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            table_name = f"{sha_type.lower()}_dictionary"
            sql = f"SELECT original FROM {table_name} WHERE encrypted = %s"
            cursor.execute(sql, (hash_str,))
            result = cursor.fetchone()
            if result is not None:  # 检查 result 是否为 None
                original_text = result['original']
                if salt:
                    original_text = original_text.replace(salt, '')
                save_to_history(hash_str, original_text, sha_type.upper())
                return original_text
            return None
    finally:
        conn.close()


# URL 编码解码
def url_encrypt(text):
    encrypted = urllib.parse.quote(text)
    save_to_history(text, encrypted, 'URL')
    return encrypted


def url_decrypt(encrypted_text):
    decrypted = urllib.parse.unquote(encrypted_text)
    save_to_history(encrypted_text, decrypted, 'URL')
    return decrypted


# 栅栏加密解密
def rail_fence_encrypt(text, key):
    fence = [['\n' for i in range(len(text))] for j in range(key)]
    dir_down = False
    row, col = 0, 0

    for i in range(len(text)):
        if (row == 0) or (row == key - 1):
            dir_down = not dir_down
        fence[row][col] = text[i]
        col += 1
        if dir_down:
            row += 1
        else:
            row -= 1

    result = []
    for i in range(key):
        for j in range(len(text)):
            if fence[i][j] != '\n':
                result.append(fence[i][j])
    encrypted = ''.join(result)
    save_to_history(text, encrypted, 'RailFence')
    save_to_dictionary(encrypted, text, 'RailFence')
    return encrypted


def rail_fence_decrypt(encrypted_text, key):
    fence = [['\n' for i in range(len(encrypted_text))] for j in range(key)]
    dir_down = None
    row, col = 0, 0

    for i in range(len(encrypted_text)):
        if row == 0:
            dir_down = True
        if row == key - 1:
            dir_down = False
        fence[row][col] = '*'
        col += 1
        if dir_down:
            row += 1
        else:
            row -= 1

    index = 0
    for i in range(key):
        for j in range(len(encrypted_text)):
            if ((fence[i][j] == '*') and (index < len(encrypted_text))):
                fence[i][j] = encrypted_text[index]
                index += 1

    result = []
    row, col = 0, 0
    for i in range(len(encrypted_text)):
        if row == 0:
            dir_down = True
        if row == key - 1:
            dir_down = False
        if fence[row][col] != '*':
            result.append(fence[row][col])
            col += 1
        if dir_down:
            row += 1
        else:
            row -= 1
    decrypted = ''.join(result)
    save_to_history(encrypted_text, decrypted, 'RailFence')
    return decrypted
